/*
* Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
* Description: hccl comm 层rpc接口实现(CPU和NPU都包含)
* Author: baochangjin
* Create: 2024-02-25
*/

#include <atomic>
#include <algorithm>
#include <arpa/inet.h>
#include <fstream>
#include <fcntl.h>
#include <unistd.h>
#include <hccl/hccl_types.h>
#include "hccl_impl_base.h"
#include "network/hccp.h"
#include "externalinput_pub.h"
#include "topoinfo_ranktableParser_pub.h"
#include "alg_remote_lookup.h"
#include "hdds_param.h"
#include "device_capacity.h"
#include "transport_heterog_roce_pub.h"
#include "hccl_impl_rpc.h"
#include "hccl_comm_pub.h"
#include "hccl_comm_rpc.h"
namespace hccl {
HcclCommRpc::HcclCommRpc(u64 inCCLbufferSize, u64 outCCLbufferSize, std::string identifier)
    : HcclCommBase(inCCLbufferSize, outCCLbufferSize, identifier)
{}

HcclCommRpc::~HcclCommRpc()
{
    impl_ = nullptr;
}

HcclResult HcclCommRpc::InitImpl(DevType deviceType)
{
    impl_.reset((new (std::nothrow) HcclImplRpc()));
    if (!isHeterogComm_) {
        deviceType_ = deviceType;
    }

    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::CreateGroup(const std::string &group, const u32 &groupRank, const u32 &userRank,
    const std::vector<u32> &groupRanks, std::shared_ptr<HcclCommBase> &groupComm)
{
    // 增加输出日志关键字
    HCCL_INFO("HCCL_KEY_INFO: group[%s], groupRank[%u], userRank[%u], groupComm[%p]", group.c_str(),
        groupRank, userRank, groupComm.get());

    // 入参有消息校验
    if (group.length() == 0) {
        HCCL_ERROR("[Create][Group]errNo[0x%016llx] group name lenth is 0", HCCL_ERROR_CODE(HCCL_E_PARA));
        return HCCL_E_PARA;
    }

    if (groupRank >= groupRanks.size()) {
        HCCL_ERROR("[Create][Group]errNo[0x%016llx] group rank[%u] out of range [0,%llu])",
            HCCL_ERROR_CODE(HCCL_E_PARA), groupRank, groupRanks.size() - 1);
        return HCCL_E_PARA;
    }

    HcclRootInfo id;
    CHK_RET(GetUniqueId(&id));

    HcclCommParams params;
    params.rank = groupRank;
    params.userRank = userRank;
    params.totalRanks = groupRanks.size();
    params.isHeterogComm = isHeterogComm_;
    s32 iret = snprintf_s(params.id.internal, HCCL_ROOT_INFO_BYTES, HCCL_ROOT_INFO_BYTES - 1, "%s%s%s",
        id.internal, "-", group.c_str());

    CHK_PRT_RET((iret == -1), HCCL_ERROR("[Create][Group]errNo[0x%016llx] get group unique id falied",
        HCCL_ERROR_CODE(HCCL_E_INTERNAL)), HCCL_E_INTERNAL);

    WorldGroupInfo groupCommonData;

    CHK_RET(impl_->GetGroupCommonData(groupCommonData));
    params.logicDevId = groupCommonData.deviceLogicId;
    params.profilingInitiated = groupCommonData.profilingInitiated;
    params.deviceType = deviceType_;
    params.hcomGroupNicInit = impl_->GetNicInitialized();
    std::vector<RankInfo> rankList;

    CHK_RET(impl_->GetGroupRanksInfo(groupRanks, rankList));

    groupComm.reset(new (std::nothrow) hccl::HcclCommBase(0, 0, group));
    CHK_SMART_PTR_NULL(groupComm);
    CHK_RET(groupComm->init(params, rankList, groupCommonData));

    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::CreateGroup(const std::string &group, const u32 &groupRank, const u32 &userRank,
    const std::vector<u32> &groupRanks, std::shared_ptr<hcclComm> &groupComm)
{
    // 增加输出日志关键字
    HCCL_INFO("HCCL_KEY_INFO: group[%s], groupRank[%u], userRank[%u], groupComm[%p]", group.c_str(),
        groupRank, userRank, groupComm.get());

    // 入参有消息校验
    if (group.length() == 0) {
        HCCL_ERROR("[Create][Group]errNo[0x%016llx] group name lenth is 0", HCCL_ERROR_CODE(HCCL_E_PARA));
        return HCCL_E_PARA;
    }

    if (groupRank >= groupRanks.size()) {
        HCCL_ERROR("[Create][Group]errNo[0x%016llx] group rank[%u] out of range [0,%llu])",
            HCCL_ERROR_CODE(HCCL_E_PARA), groupRank, groupRanks.size() - 1);
        return HCCL_E_PARA;
    }

    HcclRootInfo id;
    CHK_RET(GetUniqueId(&id));

    HcclCommParams params;
    params.rank = groupRank;
    params.userRank = userRank;
    params.totalRanks = groupRanks.size();
    params.isHeterogComm = isHeterogComm_;
    s32 iret = snprintf_s(params.id.internal, HCCL_ROOT_INFO_BYTES, HCCL_ROOT_INFO_BYTES - 1, "%s%s%s",
                          id.internal, "-", group.c_str());

    CHK_PRT_RET((iret == -1), HCCL_ERROR("[Create][Group]errNo[0x%016llx] get group unique id falied",
        HCCL_ERROR_CODE(HCCL_E_INTERNAL)), HCCL_E_INTERNAL);

    WorldGroupInfo groupCommonData;

    CHK_RET(impl_->GetGroupCommonData(groupCommonData));
    params.logicDevId = groupCommonData.deviceLogicId;
    params.profilingInitiated = groupCommonData.profilingInitiated;
    params.deviceType = deviceType_;
    params.hcomGroupNicInit = impl_->GetNicInitialized();
    std::vector<RankInfo> rankList;

    CHK_RET(impl_->GetGroupRanksInfo(groupRanks, rankList));

    groupComm.reset(new (std::nothrow) hccl::hcclComm(0, 0, group));
    CHK_SMART_PTR_NULL(groupComm);
    CHK_RET(groupComm->init(params, rankList, groupCommonData));

    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::CollRecvUpdateRequest(void* keys, int32_t keyCount, HcclDataType keyType, void* values,
    int32_t valueCount, HcclDataType valueType, int32_t tag,  ServiceHandle* handle, UpdateReqStatus* status)
{
    bool isWorker = false;
    bool isKernelLaunch = false;

    HcclOperator<EmbeddingServiceParam> *ptr = EsOpNew();
    CHK_PTR_NULL(ptr);
    *handle = ptr;
    EmbeddingServiceParam &param = ptr->param_;
    param.opType = EmbeddingServiceParam::OperatorType::RECV_UPDATE_REQUEST;
    param.tag = tag;
    param.keys = keys;
    param.keyMaxNum = keyCount;
    param.keyType = keyType;
    param.values = values;
    param.valueCount = valueCount;
    param.valuesType = valueType;
    param.getRequestParam.updateReqStatus = status;
    param.getRequestParam.commHandle = this;

    ReqStatus reqStatus{};
    param.getRequestParam.reqStatus = &reqStatus;

    CHK_RET(impl_->RegEsPsTransferMem(*ptr));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteUpdateExecutor(tag, isWorker, isKernelLaunch));

    CHK_RET(impl_->GetWorkerList(param.getRequestParam.workerList));
    CHK_RET(impl_->GetPsMap(param.getRequestParam.psMap));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->CollRecvUpdateRequest(ptr));

    status->tableId = reqStatus.tableId;
    status->globalStep = reqStatus.globalStep;
    param.getRequestParam.reqStatus = nullptr;
    param.tableId = reqStatus.tableId;

    if (param.getRequestParam.updateReqStatus->actualKeyCount == 0) {
        HcclRequest request = nullptr;
        CHK_RET(IsendUpdateResponse(*handle, &request));
    }
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::IsendUpdateResponse(ServiceHandle handle, HcclRequest* request)
{
    HcclOpPtr ptr = static_cast<HcclOpPtr>(handle);
    HcclOperator<EmbeddingServiceParam> *paraPtr = static_cast<HcclOperator<EmbeddingServiceParam> *>(handle);
    paraPtr->param_.opType = EmbeddingServiceParam::OperatorType::SEND_UPDATE_RESPONSE;
    paraPtr->param_.getRequestParam.outRequest = request;
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->IsendUpdateResponse(ptr));
    EsOpDelete(paraPtr);
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::SetKeyNumAndValCount(HcclOperator<EmbeddingServiceParam> &para, u64 keyNum, u64 keyMaxNum,
    u64 valueCount)
{
    HCCL_DEBUG("keyMaxNum[%llu] keyNum[%llu] valueCount[%llu]", keyMaxNum, keyNum, valueCount);

    // 暂时规避ge代码未合入问题
    if (keyNum == 0 || keyNum >= keyMaxNum) {
        HCCL_WARNING("keyNum[%llu] >= keyMaxNum[%llu], or keyNum is 0, use default keyMaxNum", keyNum, keyMaxNum);
        return HCCL_SUCCESS;
    }

    para.param_.keyMaxNum = keyNum;
    para.param_.valueCount = valueCount;

    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::CollRecvLookupRequest(void* keys, u32 keyCount, HcclDataType keyType,
    s32 tag, ServiceHandle* handle, LookupReqStatus* status, bool isWaitAny)
{
    HcclOperator<EmbeddingServiceParam> *ptr = EsOpNew();
    CHK_PTR_NULL(ptr);
    *handle = ptr;
    EmbeddingServiceParam &param = ptr->param_;
    param.opType = EmbeddingServiceParam::OperatorType::GET_LOOK_UP_REQUEST;
    param.tag = tag;
    param.keys = keys;
    param.keyMaxNum = keyCount;
    param.keyType = keyType;
    param.getRequestParam.lookupReqStatus = status;
    param.getRequestParam.commHandle = this;
    param.waitAny = isWaitAny;

    ReqStatus reqStatus{};

    param.getRequestParam.reqStatus = &reqStatus;

    CHK_RET(impl_->GetTagWorkerList(tag, param.getRequestParam.workerList));
    u32 workerSize = param.getRequestParam.workerList.size();

    auto it = lookupRecvRequestCounter_.Emplace(tag, 0);
    u32 &count = it.first->second;
    if (!it.second) {
        u32 &count = it.first->second;
        if (workerSize == 0) {
            HCCL_ERROR("Divisor is 0, the calculation result will be incorrect.");
            return HCCL_E_PARA;
        } else {
            count = (count + 1) % workerSize;
        }
    }

    param.getRequestParam.lookupRecvRequestCount = count;
    bool isWorker = false;
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteLookUpExecutor(tag, isWorker, false, ptr));
    CHK_RET(impl_->GetPsMap(param.getRequestParam.psMap));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->CollRecvLookupRequest(ptr));

    param.getRequestParam.actualKeyCount = reqStatus.actualSize;
    status->actualCount = reqStatus.actualSize;
    status->tableId = reqStatus.tableId;
    status->tag = reqStatus.tag;
    status->rsvd0 = reqStatus.rsvd0;
    param.getRequestParam.reqStatus = nullptr;
    param.tableId = reqStatus.tableId;

    if (status->actualCount == 0) {
        HcclRequest request = nullptr;
        void* values = nullptr;
        u32 count = 0;
        HcclDataType type = HcclDataType::HCCL_DATA_TYPE_INT8;
        CHK_RET(IsendLookupResponse(values, count, type, *handle, &request));
    }
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::IsendLookupResponse(void* values, u32 count, HcclDataType type, ServiceHandle handle,
    HcclRequest* request)
{
    HcclOpPtr ptr = static_cast<HcclOpPtr>(handle);
    HcclOperator<EmbeddingServiceParam> *paraPtr = static_cast<HcclOperator<EmbeddingServiceParam> *>(handle);
    paraPtr->param_.opType = EmbeddingServiceParam::OperatorType::SET_LOOK_UP_RESPONSE;
    paraPtr->param_.valueCount = count;
    paraPtr->param_.values = values;
    paraPtr->param_.valuesType = type;
    paraPtr->param_.getRequestParam.outRequest = request;
#ifdef ES_DATA_DFX
    u32 valueDim = count / paraPtr->param_.getRequestParam.actualKeyCount;
    std::string printInfo = "ps lookup response";
    EsUtils::PrintEsKvDataSummary(printInfo, static_cast<float *>(values),
        paraPtr->param_.getRequestParam.actualKeyCount, valueDim, static_cast<s64 *>(paraPtr->param_.keys));
#endif
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->CollIsendLookupResponse(ptr));
    EsOpDelete(paraPtr);

    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::GetLookupRequest(s32 tag, ServiceHandle* handle, u64* keys, u64 keyMaxNum, ReqStatus* request)
{
    bool isWorker = false;
    std::vector<u32> workerList;
    CHK_RET(impl_->GetTagWorkerList(tag, workerList));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteLookUpExecutor(tag, isWorker, false));

    *request = {};
    HcclOperator<EmbeddingServiceParam> *ptr = EsOpNew();
    CHK_PTR_NULL(ptr);
    *handle = ptr;
    EmbeddingServiceParam &param = ptr->param_;
    param.opType = EmbeddingServiceParam::OperatorType::GET_LOOK_UP_REQUEST;
    param.tag = tag;
    param.keys = keys;
    param.keyMaxNum = keyMaxNum;
    param.getRequestParam.workerList = workerList;
    param.getRequestParam.reqStatus = request;
    param.getRequestParam.commHandle = this;
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->GetLookupRequest(ptr));
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::IsetLookupResponse(void* values, u32 count, HcclDataType type, ServiceHandle handle,
    HcclRequest *request)
{
    HcclOpPtr ptr = static_cast<HcclOpPtr>(handle);
    HcclOperator<EmbeddingServiceParam> *paraPtr = static_cast<HcclOperator<EmbeddingServiceParam> *>(handle);
    paraPtr->param_.opType = EmbeddingServiceParam::OperatorType::SET_LOOK_UP_RESPONSE;
    paraPtr->param_.valueCount = count;
    paraPtr->param_.values = values;
    paraPtr->param_.valuesType = type;
    paraPtr->param_.getRequestParam.outRequest = request;
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->IsetLookupResponse(ptr));
    EsOpDelete(paraPtr);
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::WaitLookupData(s32 tag, void *value, void *stream, s32 &outZerocpyFlag, HcclRequest request)
{
    HcclOpPtr ptr = static_cast<HcclOpPtr>(request);
    HcclOperator<EmbeddingServiceParam> *paraPtr = static_cast<HcclOperator<EmbeddingServiceParam> *>(request);
    CHK_PTR_NULL(paraPtr);
    paraPtr->param_.opType = EmbeddingServiceParam::OperatorType::WAIT_LOOK_UP_DATA;
    paraPtr->param_.values = value;
    paraPtr->param_.embeddingParam.stream = stream;
    paraPtr->param_.outZerocpyFlag = outZerocpyFlag;

    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->WaitLookUpData(ptr));

    EsOpDelete(paraPtr);
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::RemoteLookUp(s32 tag, s32 insertFlag, void *stream, s32 &intZerocpyFlag,
    s32 &outZerocpyFlag, RemoteLookupPubParams &lookupPubInfo, HcclRequest* request, u64 keyNum)
{
    std::map<u32, u32> psMap;
    bool isWorker = true;
    CHK_RET(impl_->GetPsMap(psMap));
    // 如果是zerocopy得需要注册全局共享内存
    if (!GetExternalInputHcclIsTcpMode()) {
        CHK_RET(impl_->RegisterShareMemory());
    }
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteLookUpExecutor(tag, isWorker, false));

    HcclOperator<EmbeddingServiceParam> *ptr = EsOpNew();
    CHK_PTR_NULL(ptr);
    *request = ptr;
    EmbeddingServiceParam &param = ptr->param_;
    param.opType = EmbeddingServiceParam::OperatorType::REMOTE_LOOK_UP;
    param.keyMaxNum = lookupPubInfo.keyMaxNum;
    param.keys = lookupPubInfo.keys;
    param.keyType = HCCL_DATA_TYPE_UINT64;
    param.tag = tag;
    param.values = lookupPubInfo.value;
    param.valuesType = HCCL_DATA_TYPE_FP32;
    param.intZerocpyFlag = intZerocpyFlag;
    param.outZerocpyFlag = outZerocpyFlag;
    param.tableIdAddr = lookupPubInfo.tableId;
    param.embeddingParam.insertFlag = insertFlag;
    param.embeddingParam.valueItemSize = lookupPubInfo.valueItemSize;
    param.embeddingParam.stream = stream;
    param.embeddingParam.psMap = psMap;
    param.maxEmbeddingDim = lookupPubInfo.maxEmbeddingDim;

    CHK_RET(SetKeyNumAndValCount(*ptr, keyNum, lookupPubInfo.keyMaxNum,
        ((keyNum * lookupPubInfo.valueItemSize) / sizeof(float))));

    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RemoteLookUp(ptr));

    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::CollRemoteUpdate(s32 tag, void* keys, u32 keyMaxNum, void *values,
    void *tableId, s32 insertOption, s32 valueItemSize, rtStream_t stream, s32 maxEmbeddingDim, u64 keyNum)
{
    bool isWorker = true;
    bool isKernelLaunch = true;
    HcclOperator<EmbeddingServiceParam> para{};
    para.param_.opType = EmbeddingServiceParam::OperatorType::COLL_REMOTE_UPDATE;
    para.param_.tag = tag;
    para.param_.keys = keys;
    para.param_.keyType = HCCL_DATA_TYPE_UINT64;
    para.param_.keyMaxNum = keyMaxNum;
    para.param_.valueCount = ((static_cast<s64>(keyMaxNum) * valueItemSize) / sizeof(float));
    para.param_.values = values;
    para.param_.valuesType = HCCL_DATA_TYPE_FP32;
    para.param_.tableIdAddr = tableId;
    para.param_.embeddingParam.insertFlag = insertOption;
    para.param_.embeddingParam.valueItemSize = valueItemSize;
    para.param_.embeddingParam.stream = stream;
    para.param_.pairedMode = false;
    para.param_.disableUnique = false;
    para.param_.maxEmbeddingDim = maxEmbeddingDim;

    CHK_RET(impl_->CreateOrUpdateTransferMem(&para));

    CHK_RET(SetKeyNumAndValCount(para, keyNum, keyMaxNum, ((keyNum * valueItemSize) / sizeof(float))));

    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteUpdateExecutor(tag, isWorker, isKernelLaunch));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->CollRemoteUpdate(&para));

    return HCCL_SUCCESS;
}

void HcclCommRpc::CollRemotePairedParaAssign(HcclOperator<EmbeddingServiceParam> &para,
    const HcomRemoteOperationParams &params)
{
    para.param_.tag = params.tag;
    para.param_.sTag = std::to_string(params.tag);
    para.param_.keys = params.keyAddr;
    para.param_.keyType = HCCL_DATA_TYPE_UINT64;
    para.param_.keyMaxNum = params.count;
    para.param_.valueCount = params.count * params.valueDim;
    para.param_.values = params.value;
    para.param_.valuesType = HCCL_DATA_TYPE_FP32;
    para.param_.tableIdAddr = params.tableId;
    para.param_.embeddingParam.outputTableId = params.tableId;
    para.param_.embeddingParam.insertFlag = params.insertOption;
    para.param_.embeddingParam.valueItemSize = params.valueDim * SIZE_TABLE[para.param_.valuesType];
    para.param_.embeddingParam.stream = params.stream;
    para.param_.indices = params.indices;
    para.param_.numUniqued = params.numUniqued;
    para.param_.psSeg = params.psSeg;
    para.param_.psSegNum = params.psSegNum;
    para.param_.embeddingParam.flags = params.flags;
    para.param_.disableUnique = params.disableUnique;
    para.param_.intZerocpyFlag = params.intZerocpyFlag;
    para.param_.outZerocpyFlag = params.outZerocpyFlag;
    para.param_.globalStepAddr = params.globalStepAddr;

    para.param_.keyNumInput = params.keyNumInput;
    para.param_.uniqueIndices = params.uniqueIndices;
    para.param_.uniqued = params.uniqued;
    para.param_.keyCount = params.keyCount;
    para.param_.maxEmbeddingDim = params.maxEmbeddingDim > params.valueDim ? params.maxEmbeddingDim : params.valueDim;
    para.param_.uniqueTag = params.uniqueTag;
}

HcclResult HcclCommRpc::CollRemoteUpdatePaired(const HcomRemoteOperationParams &params)
{
    bool isWorker = true;
    bool isKernelLaunch = true;
    HcclOperator<EmbeddingServiceParam> para{};

    CollRemotePairedParaAssign(para, params);
    para.param_.opType = EmbeddingServiceParam::OperatorType::COLL_REMOTE_UPDATE_PAIRED;
    para.param_.pairedMode = true;
    para.param_.disableUnique = false;

    CHK_RET(impl_->CreateOrUpdateTransferMem(&para));

    CHK_RET(SetKeyNumAndValCount(para, params.keyNum, params.count, params.keyNum * params.valueDim));

    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteUpdateExecutor(params.tag, isWorker,
        isKernelLaunch));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->CollRemoteUpdatePaired(&para));

    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::CollRemoteLookup(s32 tag, void* keys, u32 keyMaxNum, void *values,
    void *tableId, s32 insertOption, s32 valueItemSize, rtStream_t stream, s32 flags, s32 maxEmbeddingDim, u64 keyNum)
{
    bool isWorker = true;
    HcclOperator<EmbeddingServiceParam> para{};
    para.param_.opType = EmbeddingServiceParam::OperatorType::COLL_REMOTE_LOOK_UP;
    para.param_.tag = tag;
    para.param_.sTag = std::to_string(tag);
    para.param_.keys = keys;
    para.param_.keyType = HCCL_DATA_TYPE_UINT64;
    para.param_.keyMaxNum = keyMaxNum;
    para.param_.valueCount = ((keyMaxNum * valueItemSize) / sizeof(float));
    para.param_.values = values;
    para.param_.valuesType = HCCL_DATA_TYPE_FP32;
    para.param_.embeddingParam.outputTableId = tableId;
    para.param_.tableIdAddr = tableId;
    para.param_.embeddingParam.insertFlag = insertOption;
    para.param_.embeddingParam.valueItemSize = valueItemSize;
    para.param_.embeddingParam.stream = stream;
    para.param_.embeddingParam.flags = flags;
    para.param_.pairedMode = false;
    para.param_.disableUnique = false;
    para.param_.intZerocpyFlag = false;
    para.param_.outZerocpyFlag = false;
    para.param_.uniqued = false;
    para.param_.maxEmbeddingDim = maxEmbeddingDim;

    // 中转内存、stream、notify申请
    CHK_RET(impl_->RegEsWorkSpaceResource(&para));

    CHK_RET(SetKeyNumAndValCount(para, keyNum, keyMaxNum, ((keyNum * valueItemSize) / sizeof(float))));

    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteLookUpExecutor(tag, isWorker, true));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->CollRemoteLookUp(&para));
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::CollRemoteLookUp(const HcomRemoteOperationParams &params)
{
    bool isWorker = true;
    HcclOperator<EmbeddingServiceParam> para{};
    para.param_.opType = EmbeddingServiceParam::OperatorType::COLL_REMOTE_LOOK_UP;
    CollRemotePairedParaAssign(para, params);
    para.param_.pairedMode = false;

    // 中转内存、stream、notify申请
    CHK_RET(impl_->RegEsWorkSpaceResource(&para));

    CHK_RET(SetKeyNumAndValCount(para, params.keyNum, params.count, params.keyNum * params.valueDim));

    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteLookUpExecutor(params.tag, isWorker, true));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->CollRemoteLookUp(&para));
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::CollRemoteLookupPaired(const HcomRemoteOperationParams &params)
{
    bool isWorker = true;
    HcclOperator<EmbeddingServiceParam> para{};

    CollRemotePairedParaAssign(para, params);
    para.param_.opType = EmbeddingServiceParam::OperatorType::COLL_REMOTE_LOOK_UP_PAIRED;
    para.param_.pairedMode = true;
    para.param_.disableUnique = false;
    para.param_.uniqued = false;

    CHK_RET(impl_->RegEsWorkSpaceResource(&para));

    CHK_RET(SetKeyNumAndValCount(para, params.keyNum, params.count, params.keyNum * params.valueDim));

    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteLookUpExecutor(params.tag, isWorker, true));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->CollRemoteLookUpPaired(&para));
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::CollRemoteLookupUniquedPaired(const HcomRemoteOperationParams &params)
{
    bool isWorker = true;
    HcclOperator<EmbeddingServiceParam> para{};

    CollRemotePairedParaAssign(para, params);
    para.param_.opType = EmbeddingServiceParam::OperatorType::COLL_REMOTE_LOOK_UP_PAIRED;
    para.param_.pairedMode = true;
    para.param_.disableUnique = false;
    para.param_.uniqued = true;

    CHK_RET(impl_->RegEsWorkSpaceResource(&para));

    CHK_RET(SetKeyNumAndValCount(para, params.keyNum, params.count, params.keyNum * params.valueDim));

    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegRemoteLookUpExecutor(params.tag, isWorker, true));
    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->CollRemoteLookUpPaired(&para));
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::AbortSelf(s32 tag)
{
    return reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->AbortSelf(tag);
}

HcclResult HcclCommRpc::RegTransportLinks(s32 linkNum, void *transportPara)
{
    return reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->RegTransportLinks(linkNum, transportPara);
}

HcclResult HcclCommRpc::PsGetBoardVersion(bool &isBoardVersion)
{
    isBoardVersion = reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->PsGetBoardVersion();

    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::WorkerGetBoardVersion(bool &isBoardVersion)
{
    isBoardVersion = reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->WorkerGetBoardVersion();

    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::WaitSome(s32 count, HcclRequest requestArray[], s32* compCount, s32 compIndices[],
    HcclStatus compStatus[])
{
    HcclOperator<EmbeddingServiceParam> op;
    op.param_.opType = EmbeddingServiceParam::OperatorType::WAIT_SOME;
    op.param_.waitSomeParam.requestCount = count;
    op.param_.waitSomeParam.requestArray = requestArray;
    op.param_.waitSomeParam.compCount = compCount;
    op.param_.waitSomeParam.compIndices = compIndices;
    op.param_.waitSomeParam.compStatus = compStatus;

    CHK_RET(reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->WaitSome(&op));
    return HCCL_SUCCESS;
}

HcclResult HcclCommRpc::Isend(void *buffer, u64 count, HcclDataType dataType, u32 peerRank,
    s32 tag, HcclRequest &request, HcclUserRequire &userRequire) const
{
    /* 入参检查 */
    CHK_RET(impl_->CheckCount(count));
    CHK_RET(impl_->CheckDataType(dataType, false));
    return reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->
        Isend(buffer, count, dataType, peerRank, tag, request, userRequire);
}

HcclResult HcclCommRpc::Improbe(u32 peerRank, s32 tag, s32 &flag, HcclMessage &msgHandle,
    HcclStatus &status) const
{
    CHK_SMART_PTR_NULL(impl_);
    return reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->Improbe(peerRank, tag, flag, msgHandle, status);
}

HcclResult HcclCommRpc::Imrecv(void *buffer, u64 count, HcclDataType dataType,
    HcclMessage msg, HcclRequest &request)
{
    CHK_SMART_PTR_NULL(impl_);

    /* 入参检查 */
    CHK_RET(impl_->CheckCount(count));
    CHK_RET(impl_->CheckDataType(dataType, false));
    return reinterpret_cast<hccl::HcclImplRpc *>(impl_.get())->Imrecv(buffer, count, dataType, msg, request);
}
}